#!/usr/bin/perl

use strict;
use Sys::Syslog;
use Getopt::Std;

my %config = (
		'facility',			"news",
		'detach',			"true",
		'file',				"/etc/hosts.deny",
		'pidfile',			"/var/run/nnrpdlimit.pid",
		'expire',			86400,
		'frequence',			600,
		"max_connection_time", 		10800,
		'error_too_long_connection',	"Persistent connections are forbidden",
		'max_cpu_time',			6,		# CPU time in seconds like ps
		'error_too_much_cpu_time',      "Too much CPU time",
		'max_concurrent_connections',   2,
		'error_too_many_concurrent_connections', "Too many concurrent connections",
		'max_memory_usage',		40000, # Kb
		'error_too_much_ram',		"Too many queries (RAM)"
	     );

my $proc = "/proc";

my $pid;

my %options = ();
getopts("f:t:e:p:g:a:r:s:q:DZh",\%options);

$config{'file'} 			=  $options{f} if (defined $options{f});
$config{'frequence'} 			=  $options{t} if (defined $options{t});
$config{'expire'} 			=  $options{e} if (defined $options{e});
$config{'max_connection_time'} 		=  $options{p} if (defined $options{p});
$config{'max_cpu_time'} 		=  $options{g} if (defined $options{g});
$config{'max_concurrent_connections'} 	=  $options{a} if (defined $options{a});
$config{'max_memory_usage'}	   	=  $options{r} if (defined $options{r});
$config{'facility'}		        =  $options{s} if (defined $options{s});
$config{'pidfile'}                      =  $options{q} if (defined $options{q});
$config{'detach'} 			=  "false"     if (defined $options{D});

if (defined $options{h})
{
	print <<EOT;
nnrpdlimit v. 0.7.7 Copyright 2009 by Paolo Amoroso <freedom\@aioe.org> (BSD License)

Syntax:
	nnrpdlimit [ options ]
	nnrpdlimit -h

Options:
	-f <file>	 : Set the file to use for both input and output
			   Default: /etc/hosts.deny
	-t <seconds>	 : Make a check every <seconds> seconds. 
			   Default: 600 (10 minutes)
	-e <seconds>     : Unban banned IP addresses after <seconds>
			   Default: 86400 (1 Day)
	-p <seconds>     : Ban every client that stays connected more than <seconds> 
			   Default: 10800 (3 Hours)
	-g <seconds>	 : Ban every client that consumes more than <seconds> of CPU time
			   Default: 6
	-a <connections> : Ban every client that establishes more than <connections> per time
			   Default: 2
	-r <Kbytes>	 : Maximum amount of RAM that a client can allocate without being banned
			   Default: 40000 Kb
	-s <facility>	 : Syslog facility to use
			   Default: news
	-q <file>	 : PIDFile to use
			   Default: /var/run/nnrpdlimit.pid

	-D		 : Run in foreground
	-Z		 : Kill a running instance of nnrpdlimit
	-h		 : This help

EOT
	exit(0);
}

if (defined $options{Z})
{
	&kill_nnrpdlimit();
	exit(0);
}

&init_checks();

if ( $config{'detach'} eq "true" )
{
	$pid = fork();
} else
{
	$pid = 0;
}

if (not defined $pid) 
{
	die "resources not avilable.\n";
} elsif ($pid == 0) 
{
	while (1)
	{
		my %banned;
		my @pids_to_kill;
		my %ips;
		opendir(my $PROC, "$proc");

		if (!$PROC)
		{
			my $text = "Unable to read $proc virtual filesystem";
			print STDERR "$text\n" if ( $config{'detach'} ne "true" );
			&log( "crit", $text );
			exit 5;
		}

		my @processes = readdir($PROC);
		closedir($PROC);

		&expire_old_data();
		my @pids;

		foreach (@processes)
		{
			if ( $_ =~ /^[0-9]/ )
			{
				push ( @pids, $_ );
			}
		}

		foreach (@pids)
		{
        		my $cmdline =  &get_command_line($_);	
			push ( @pids_to_kill, $_ ) if ( $cmdline =~ /nnrpd-ssl.+\-s/ );

			if ( $cmdline =~ /nnrpd\:|nnrpd-ssl\:/ )
			{
				$cmdline =~  /([0-9|\.]+)/;
				my $ip = $1;
				if (defined $ips{$ip}) 
				{
					$ips{$ip}++;
					if ( $ips{$ip} > $config{'max_concurrent_connections'} )
					{
						&log( "notice", "Banning $ip, $ips{$ip} concurrent connections, max is  $config{'max_concurrent_connections'}" );
						$banned{$ip} = $config{'error_too_many_concurrent_connections'};
						push ( @pids_to_kill, $_ );
						$ips{$ip} = -50; # FIXME: needed in order to avoid repetitions
					}
				} 
				else
				{
					$ips{$ip} = 1;
				}
				my $active_time = &get_active_time($_);

				if ( $active_time > $config{'max_connection_time'} )
				{
					&log( "notice", "Banning $ip, $active_time connection time, max is $config{'max_connection_time'}" );
					$banned{$ip} = $config{'error_too_long_connection'};
					push ( @pids_to_kill, $_ );
				}		

				my ($success,$arg) = &check_cpu_time($_);
				if ( $success == 1 )
                        	{
					&log( "notice", "Banning $ip, too much CPU time ($arg seconds)" );
                                	$banned{$ip} = $config{'error_too_much_cpu_time'};
                                	push ( @pids_to_kill, $_ );
                        	}
				elsif ( $success == 2 )
                                {
					&log( "notice", "Banning $ip, too much RAM used ($arg KBytes)" );					
                                        $banned{$ip} = $config{'error_too_much_ram'};
                                        push ( @pids_to_kill, $_ );
                                }

		
			}
		}

        	open my $FILE, ">>$config{'file'}";

		if (!$FILE)
		{
			my $text = "Unable to append data to $config{'file'}";
			print STDERR "$text\n" if ( $config{'detach'} ne "true" );
			&log("crit", $text );
			exit 5;
		}

		my $time = time();

		foreach (keys %banned)
		{
			&log( "notice", "IP $_ banned: $banned{$_}" );
			print $FILE "nnrpd:$_: twist echo -e \"502 $banned{$_}\\r\\n\" # $time\n" if ( $_ ne "" );
		}

		close $FILE;

		foreach (@pids_to_kill)
		{
			&log( "debug", "Process $_ killed" );
			system("kill $_");
		}

		&log( "notice", "Data checked" );
		sleep $config{'frequence'};
	}
} else 
{
	exit(0);
}

sub check_cpu_time($)
{
	my $file = "/proc/$_[0]/stat";
	open my $PROC, "$file";

        if (!$PROC)
        {
		my $text = "Unable to read $file";
		print STDERR "$text\n" if ( $config{'detach'} ne "true" );
	        &log( "crit", $text);
                exit 5;
        }

	my $line = <$PROC>;
	close $PROC;
	my @elems = split ( /\ /, $line );
	my $tot = $elems[13] + $elems[14];
	$tot = sprintf("%d", $tot / 100);
	return (1,$tot) if ( $tot > $config{'max_cpu_time'});

	my $ram = sprintf( "%d", ($elems[22] / 1024));
	return (2,$ram) if ( $ram > $config{'max_memory_usage'} );

	return (0,0);
}

sub expire_old_data()
{
	my $file = $config{'file'};
	open my $FILE, "$file";

        if (!$FILE)
        {
		my $text = "Unable to read data from $file";
		print STDERR "$text\n" if ( $config{'detach'} ne "true" );
	        &log("crit", $text);
                exit 5;
        }

	my @lines = <$FILE>;
	close $FILE;

	open $FILE, ">$file";

        if (!$FILE)
        {
                my $text = "Unable to write data to $file";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log("crit", $text);
             	exit 5;
        }

	foreach (@lines)
	{
		if ( $_  =~ /^nnrpd/ )
		{
			$_ =~ /\#\ ([0-9]+)/;
			my $set_time = $1;
			my $now = time();
			if ( ($now - $set_time) <= $config{'expire'})
			{
	                        print $FILE $_;
			}
			else
			{
				$_  =~ /^nnrpd\:(.+)\:/;
				my $ip = $1;
				&log( "notice", "$ip unbanned after $config{'expire'} seconds" );
			}
		}
		else
		{
			print $FILE $_;
		}
	}
	
	close $FILE;

}

sub log()    # log( severity, string1, string2.... );
{
        my @strings = @_;
        openlog ( 'nnrpdlimit', 'cons,pid', $config{'facility'} );

        my $severity = shift(@strings);

        foreach ( @strings )
        {
                syslog( $severity, $_ );
        }

        closelog();
}


sub init_checks()
{
	my $name = `whoami`;
	$name =~ s/\n//;

	if ( $name ne "root" )
	{
		my $text = "nnrpdlimit must be executed by root not by $name";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "crit", $text);
                exit 5;
	}

	my $tf = $config{'pidfile'};

	my $GPP;
	open $GPP, "$tf";

	if ($GPP)
	{
		my $pid = <$GPP>;
		close $GPP;
		$pid =~ s/\n//;
		if ( $pid ne "" )
		{
			my $text = "Found an existing pidfile ($config{'pidfile'}) with $pid pid";
                	print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                	&log( "err", $text);		

			my $cmdline = &get_command_line($pid);
                 	if ( $cmdline =~ /nnrpdlimit/ )
		 	{
				my $text = "Found another working instance (pid $pid) of nnrpdlimit";
	                	print STDERR "$text\n" if ( $config{'detach'} ne "true" );
        	        	&log( "crit", $text);
				exit 5;
		 	}
		}
	}

	my $version = "/proc/version";


	my $PROC;
	open $PROC, "$version";
	if (!$PROC)
	{
	        my $text = "Unable to read $version";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "crit", $text);
                exit 5;
	}

	my $line = <$PROC>;
	close $PROC;

	if ( $line !~ /^Linux/ )
	{
		my $text = "nnrpdlimit works only with Linux, other Operative Systems aren't supported";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "crit", $text);
                exit 5;
	}

	my $cpid = $$ + 2;

	my $PID;
	open $PID, ">$tf";

	if (!$PID)
	{
                my $text = "Unable to write pidfile at $tf";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "crit", $text);
                exit 5;
	}	

	print $PID "$cpid\n";
	close $PID;
}

sub kill_nnrpdlimit()
{
	my $tf = $config{'pidfile'};

        my $GPP;
        open $GPP, "$tf";	
	my $pid = <$GPP>;	
	$pid =~ s/\n//;
        if ( $pid eq "" )
        {
        	my $text = "No other nnrpdlimit instances are running";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "err", $text);
		return 0;
	}
	else
	{
		close $GPP;
		my $status =  `kill -9 $pid`;
		$status =  `rm $tf`;	
	}
	return 0;
}

sub get_command_line($)
{
	my $pid = $_[0];
	my $file = "/proc/$pid/cmdline";
	open my $CMDLINE, "$file";
	return undef if (!$CMDLINE);
	my $cmdline = <$CMDLINE>;
	close $CMDLINE;
	return $cmdline;
}

sub get_active_time($)
{
	my $pid = $_[0];
	my $file = "/proc/$_[0]/stat";
	open my $PROC, "$file";

        if (!$PROC)
        {
                my $text = "Unable to read $file";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "crit", $text);
                exit 5;
        }

	my $line = <$PROC>;
	close $PROC;

	my @elems = split( /\ /, $line );


	$file = "/proc/stat";
	open $PROC, "$file";

        if (!$PROC)
        {
                my $text = "Unable to read $file";
                print STDERR "$text\n" if ( $config{'detach'} ne "true" );
                &log( "crit", $text);
                exit 5;
        }

	my @proc = <$PROC>;
	close $PROC;

	my $bootime;

	foreach ( @proc )
	{
		if ( $_ =~ /btime/ )
		{
			$_  =~ /^btime\ ([0-9]+)/;
			$bootime = $1;		
		}
	}

	my $due = sprintf("%d", ($elems[21]/100) );
	my $tre = $bootime + $due;
	my $dur = time() - $tre;
	return $dur;
}
